#include <iostream>
#include <onnxruntime_cxx_api.h>
#include <cpu_provider_factory.h>
#include "hcn.h"

#define debug

namespace myHcn {

static float input_data2[4][1][1][96] = {0};
static std::vector<int64_t> input_shape1 = {1,};
static std::vector<int64_t> input_shape2 = {4, 1, 1, 96};
static Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "hcn");
static Ort::SessionOptions session_options{};

void Hcn::printTensorValue(Ort::Value &value, std::string valueInfo)
{
    std::cout << "============="<< valueInfo << "================" << std::endl;
    float *floatData = value.GetTensorMutableData<float>();
    for (size_t i = 0; i < value.GetTensorTypeAndShapeInfo().GetElementCount(); i++) {
      std::cout << floatData[i] << ", ";
    }
    std::cout << std::endl;
}

void Hcn::addName(std::vector<const char*> &names, char *name)
{
    char *mem = new char(sizeof(name));
    memcpy(mem, name, sizeof(name));
    names.push_back(mem);
}

Hcn::Hcn(const char* modelFile) 
 : memory_info{Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU)},
 session{env, modelFile, session_options}
{ 
    auto input_name1 = session.GetInputNameAllocated(0,allocator);
    auto input_name2 = session.GetInputNameAllocated(1,allocator);
    addName(input_names, input_name1.get());
    addName(input_names, input_name2.get());

    auto output_name1 = session.GetOutputNameAllocated(0,allocator);
    auto output_name2 = session.GetOutputNameAllocated(1,allocator);
    addName(output_names, output_name1.get());
    addName(output_names, output_name2.get());
    run_options = Ort::RunOptions{ nullptr };
}

Hcn::~Hcn()
{
    for (auto data : input_names) {
        delete data;
    }
    for (auto data : output_names) {
        delete data;
    }
}

std::vector<Ort::Value> *Hcn::getOutputTensors()
{
    return &output_tensors;
}

Ort::Value *Hcn::run(float tfValue)
{
    input_tensors.clear();
    input_tensors.push_back(Ort::Value::CreateTensor<float>(memory_info, &tfValue, 1, input_shape1.data(), input_shape1.size()));
    if (isReRun) {
        input_tensors.push_back(Ort::Value::CreateTensor<float>(memory_info, &input_data2[0][0][0][0], 4*96, input_shape2.data(), input_shape2.size()));
        isReRun = false;
    } else {
        input_tensors.push_back(std::move(output_tensors[1]));
    }

    output_tensors.clear();
    output_tensors = session.Run(run_options, input_names.data(), input_tensors.data(), 2, output_names.data(), 2);
        
    // printTensorValue(output_tensors[0], "output_tensors[0]");
    // printTensorValue(output_tensors[1], "output_tensors[1]");
    return &output_tensors[0];
}

void Hcn::shotOverCallback()
{
    isReRun = true;
}

}   // namespace myHcn
